#include "GBufferProcess.h"
#include "Utils/LogUtil.h"
#include "Render/RenderState.h"
#include "Core/CoreGlobals.h"
#include "Game/Scene/Entity.h"
#include "Game/Component/Model.h"
#include "Render/Buffer/VertexBuffer.h"
#include "Render/Buffer/IndexBuffer.h"
#include "Render/Buffer/CustomBuffer.h"
#include "Render/Texture/CustomTexture.h"
#include "LocalShader.h"
#include "Render/VulkanTool.h"
#include "RayTracingMath.h"
#include "LocalScene.h"

void GBufferProcess::OnInit()
{
	mLocalScene = static_cast<LocalScene*>(mGameScene);

	PerFrameUbo perFrameUbo = mLocalScene->mPerFrameUbo;

	ENQUEUE_RENDER_COMMAND(BuildModel)([this, perFrameUbo](VkCommandBuffer commandBuffer)
		{
			for (uint32 i = 0; i < mLocalScene->mOpaqueEntitys.size(); ++i)
			{
				LocalEntity* opaqueEntity = mLocalScene->mOpaqueEntitys[i];
				Model& model = opaqueEntity->GetComponent<Model>();

				std::vector<VertexGpu> tempVertexs;
				tempVertexs.resize(model.mVertexs.size());
				for (uint32 j = 0; j < model.mVertexs.size(); ++j)
				{
					tempVertexs[j].position = model.mVertexs[j].Position;
					tempVertexs[j].normal = model.mVertexs[j].Normal;
					tempVertexs[j].tangent = model.mVertexs[j].Tangent;
					tempVertexs[j].uv = model.mVertexs[j].TexC;
				}

				opaqueEntity->mVertexBuffer = VertexBuffer::Builder()
					.SetData(tempVertexs.data())
					.SetSize(tempVertexs.size() * sizeof(VertexGpu))
					.SetRayUse()
					.Build();

				opaqueEntity->mIndexBuffer = IndexBuffer::Builder()
					.SetData(model.mIndices.data())
					.SetSize(model.mIndices.size() * sizeof(uint32))
					.SetIndexCount(model.mIndices.size())
					.SetRayUse()
					.Build();

				ObjectDescGpu rayObjDesc;
				rayObjDesc.vertexAddress = opaqueEntity->mVertexBuffer->GetDeviceAddress();
				rayObjDesc.indexAddress = opaqueEntity->mIndexBuffer->GetDeviceAddress();

				mObjectDescs.push_back(rayObjDesc);

				mMaterials.push_back(opaqueEntity->mMaterial);
			}

			mMaterialBuffer = CustomBuffer::Builder()
				.SetData(mMaterials.data())
				.SetSize(sizeof(MaterialGpu) * mMaterials.size())
				.SetUsage(VK_BUFFER_USAGE_TRANSFER_DST_BIT | VK_BUFFER_USAGE_STORAGE_BUFFER_BIT)
				.Build();

			mObjectDescBuffer = CustomBuffer::Builder()
				.SetData(mObjectDescs.data())
				.SetSize(sizeof(ObjectDescGpu) * mObjectDescs.size())
				.SetUsage(VK_BUFFER_USAGE_TRANSFER_DST_BIT | VK_BUFFER_USAGE_STORAGE_BUFFER_BIT)
				.Build();

			PerFrameUbo perFrameUboData = perFrameUbo;

			mPerFrameUboBuffer = CustomBuffer::Builder()
				.SetData(&perFrameUboData)
				.SetSize(sizeof(PerFrameUbo))
				.SetMemoryProperty(VK_MEMORY_PROPERTY_HOST_VISIBLE_BIT | VK_MEMORY_PROPERTY_HOST_COHERENT_BIT)
				.Build();

			mDescriptor = RHIDescriptor::Builder()
				.PushShader<GBufferVertexShader>()
				.PushShader<GBufferFragmentShader>()
				.SetDescriptorCount(0, 1)
				.Build();

			mDescriptor->SetValue("perFrameUBO", mPerFrameUboBuffer);
			mDescriptor->SetValue("Materials", mMaterialBuffer);
			mDescriptor->SetValue("objectDescs", mObjectDescBuffer);
			mDescriptorMask = mDescriptor->AllocateDescriptor();

			mGBuffer1Texture = ColorRenderTarget::Builder()
				.SetFormat(VK_FORMAT_R8G8B8A8_UNORM)
				.SetUsage(VK_IMAGE_USAGE_COLOR_ATTACHMENT_BIT | VK_IMAGE_USAGE_SAMPLED_BIT)
				.SetStoreOp(VK_ATTACHMENT_STORE_OP_STORE)
				.SetFinalLayout(VK_IMAGE_LAYOUT_SHADER_READ_ONLY_OPTIMAL)
				.Build(this, "GBuffer1Texture");

			mGBuffer2Texture = ColorRenderTarget::Builder()
				.SetFormat(VK_FORMAT_R16G16B16A16_SFLOAT)
				.SetUsage(VK_IMAGE_USAGE_COLOR_ATTACHMENT_BIT | VK_IMAGE_USAGE_SAMPLED_BIT)
				.SetStoreOp(VK_ATTACHMENT_STORE_OP_STORE)
				.SetFinalLayout(VK_IMAGE_LAYOUT_SHADER_READ_ONLY_OPTIMAL)
				.Build(this, "GBuffer2Texture");

			mGBuffer3Texture = ColorRenderTarget::Builder()
				.SetFormat(VK_FORMAT_R16G16B16A16_SFLOAT)
				.SetUsage(VK_IMAGE_USAGE_COLOR_ATTACHMENT_BIT | VK_IMAGE_USAGE_SAMPLED_BIT)
				.SetStoreOp(VK_ATTACHMENT_STORE_OP_STORE)
				.SetFinalLayout(VK_IMAGE_LAYOUT_SHADER_READ_ONLY_OPTIMAL)
				.Build(this, "GBuffer3Texture");

			mDepthRenderTarget = DepthRenderTarget::Builder()
				.Build();

			mPipeline = GraphicsPipeline::Builder()
				.SetDescriptor(mDescriptor)
				.AddColorRenderTarget(mGBuffer1Texture)
				.AddColorRenderTarget(mGBuffer2Texture)
				.AddColorRenderTarget(mGBuffer3Texture)
				.SetDepthRenderTarget(mDepthRenderTarget)
				.SetDstStageMask(VK_PIPELINE_STAGE_COLOR_ATTACHMENT_OUTPUT_BIT)
				.SetDstAccessMask(VK_ACCESS_COLOR_ATTACHMENT_WRITE_BIT)
				.SetSrcStageMask(VK_PIPELINE_STAGE_COLOR_ATTACHMENT_OUTPUT_BIT)
				.SetSrcAccessMask(VK_ACCESS_COLOR_ATTACHMENT_WRITE_BIT)
				.SetCullMode(VK_CULL_MODE_NONE)
				.Build();

			mRenderPass = RHIRenderPass::Builder()
				.AddRenderTarget(mGBuffer1Texture)
				.AddRenderTarget(mGBuffer2Texture)
				.AddRenderTarget(mGBuffer3Texture)
				.AddRenderTarget(mDepthRenderTarget)
				.AddPipeline(mPipeline)
				.SetSrcStageMask(VK_PIPELINE_STAGE_FRAGMENT_SHADER_BIT)
				.SetSrcAccessMask(VK_ACCESS_SHADER_READ_BIT)
				.SetDstStageMask(VK_PIPELINE_STAGE_FRAGMENT_SHADER_BIT)
				.SetDstAccessMask(VK_ACCESS_SHADER_READ_BIT)
				.Build();

			mFrameBuffer = RHIFrameBuffer::Builder()
				.SetRenderPass(mRenderPass)
				.Build();
		});
}

void GBufferProcess::OnTick(float deltaTime)
{
	PerFrameUbo perFrameUbo = mLocalScene->mPerFrameUbo;

	std::vector<GBufferPushConstants> gBufferPushConstants = mLocalScene->mGBufferPushConstants;

	ENQUEUE_RENDER_COMMAND(RecordCommand)([this, perFrameUbo, gBufferPushConstants](VkCommandBuffer commandBuffer)
		{
			PerFrameUbo perFrameUboData = perFrameUbo;

			mPerFrameUboBuffer->WriteMemory(&perFrameUboData);

			VkRenderPassBeginInfo renderPassBeginInfo = mFrameBuffer->GetRenderPassBeginInfo();

			vkCmdBeginRenderPass(commandBuffer, &renderPassBeginInfo, VK_SUBPASS_CONTENTS_INLINE);

			vkCmdSetViewport(commandBuffer, 0, 1, &GRenderState->GetViewport());
			vkCmdSetScissor(commandBuffer, 0, 1, &GRenderState->GetScissor());

			vkCmdBindPipeline(commandBuffer, VK_PIPELINE_BIND_POINT_GRAPHICS, mPipeline->GetPipeline());

			vkCmdBindDescriptorSets(commandBuffer, VK_PIPELINE_BIND_POINT_GRAPHICS, mPipeline->GetPipelineLayout(), 0, 1, &mDescriptorMask.GetDescriptorSet(), 0, nullptr);

			for (uint32 i = 0; i < mLocalScene->mOpaqueEntitys.size(); ++i)
			{
				LocalEntity* opaqueEntity = mLocalScene->mOpaqueEntitys[i];

				VkBuffer vertexBuffer = opaqueEntity->mVertexBuffer->GetBuffer();
				VkDeviceSize offsets = 0;
				vkCmdBindVertexBuffers(commandBuffer, 0, 1, &vertexBuffer, &offsets);

				VkBuffer indexBuffer = opaqueEntity->mIndexBuffer->GetBuffer();
				vkCmdBindIndexBuffer(commandBuffer, indexBuffer, 0, VK_INDEX_TYPE_UINT32);

				vkCmdPushConstants(commandBuffer, mPipeline->GetPipelineLayout(), VK_SHADER_STAGE_VERTEX_BIT | VK_SHADER_STAGE_FRAGMENT_BIT, 0, sizeof(GBufferPushConstants), &gBufferPushConstants[i]);

				vkCmdDrawIndexed(commandBuffer, opaqueEntity->mIndexBuffer->GetIndexCount(), 1, 0, 0, 0);
			}

			vkCmdEndRenderPass(commandBuffer);
		});
}

void GBufferProcess::OnResize()
{
	mGBuffer1Texture = ColorRenderTarget::Builder()
		.SetFormat(VK_FORMAT_R8G8B8A8_UNORM)
		.SetUsage(VK_IMAGE_USAGE_COLOR_ATTACHMENT_BIT | VK_IMAGE_USAGE_SAMPLED_BIT)
		.SetStoreOp(VK_ATTACHMENT_STORE_OP_STORE)
		.SetFinalLayout(VK_IMAGE_LAYOUT_SHADER_READ_ONLY_OPTIMAL)
		.Build(this, "GBuffer1Texture");

	mGBuffer2Texture = ColorRenderTarget::Builder()
		.SetFormat(VK_FORMAT_R16G16B16A16_SFLOAT)
		.SetUsage(VK_IMAGE_USAGE_COLOR_ATTACHMENT_BIT | VK_IMAGE_USAGE_SAMPLED_BIT)
		.SetStoreOp(VK_ATTACHMENT_STORE_OP_STORE)
		.SetFinalLayout(VK_IMAGE_LAYOUT_SHADER_READ_ONLY_OPTIMAL)
		.Build(this, "GBuffer2Texture");

	mGBuffer3Texture = ColorRenderTarget::Builder()
		.SetFormat(VK_FORMAT_R16G16B16A16_SFLOAT)
		.SetUsage(VK_IMAGE_USAGE_COLOR_ATTACHMENT_BIT | VK_IMAGE_USAGE_SAMPLED_BIT)
		.SetStoreOp(VK_ATTACHMENT_STORE_OP_STORE)
		.SetFinalLayout(VK_IMAGE_LAYOUT_SHADER_READ_ONLY_OPTIMAL)
		.Build(this, "GBuffer3Texture");

	mDepthRenderTarget = DepthRenderTarget::Builder()
		.Build();

	mFrameBuffer = RHIFrameBuffer::Builder()
		.AddRenderTarget(mGBuffer1Texture)
		.AddRenderTarget(mGBuffer2Texture)
		.AddRenderTarget(mGBuffer3Texture)
		.AddRenderTarget(mDepthRenderTarget)
		.SetRenderPass(mRenderPass)
		.Build();
}


